package main

import (
	"errors"
	"fmt"
	"gitee.com/jjawesomejj/awesome-util/config"
	"gitee.com/jjawesomejj/awesome-util/helper/pool"
	"io"
	"net"
	"sync"
	"syscall"
)

func main() {
	config.InitConfig()
	config.HandleCommandLineSetToConfig()
	listenPort := config.GetConfigByKey("listen_port").(string)
	targetAddress := config.GetConfigByKey("target_address").(string)
	server, err := net.Listen("tcp", ":"+listenPort)
	if err != nil {
		panic(err)
	}
	fmt.Println("server start at:", listenPort)
	for {
		conn, err := server.Accept()
		if err != nil {
			fmt.Println(err)
			continue
		}
		fmt.Println("new connection:", conn.RemoteAddr().String())
		go handleConn(conn, targetAddress)
	}
}

func handleConn(conn net.Conn, targetAddress string) {
	defer conn.Close()
	remoteConn, err := net.Dial("tcp", targetAddress)
	if err != nil {
		fmt.Println(err)
		return
	}
	defer remoteConn.Close()
	relay(conn, remoteConn)

}

func copyBuffer(dst io.Writer, src io.Reader) error {
	buf := pool.GetBuf(20 * 1024)
	defer pool.PutBuf(buf)
	_, err := io.CopyBuffer(dst, src, buf)
	if ne, ok := err.(net.Error); ok && ne.Timeout() {
		return nil /* ignore I/O timeout */
	} else if errors.Is(err, syscall.EPIPE) {
		return nil /* ignore broken pipe */
	} else if errors.Is(err, syscall.ECONNRESET) {
		return nil /* ignore connection reset by peer */
	}
	return err
}

func relay(left, right net.Conn) {
	wg := sync.WaitGroup{}
	wg.Add(2)
	go func() {
		defer wg.Done()
		defer func() {
			_ = left.Close()
			_ = right.Close()
		}()
		if err := copyBuffer(right, left); err != nil {
			fmt.Printf("[TCP] %v", err)
		}
	}()
	go func() {
		defer wg.Done()
		defer func() {
			_ = left.Close()
			_ = right.Close()
		}()
		if err := copyBuffer(left, right); err != nil {
			fmt.Printf("[TCP] %v", err)
		}
	}()
	wg.Wait()
}
